import tensorflow as tf
import os
import matplotlib.pyplot as plt

if __name__ == '__main__':
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
    # 准备 x 输入向量
    x = tf.linspace(-10, 10, 100)

    y = tf.nn.leaky_relu(x, alpha=0.1)
    # print(y)

    x = tf.Variable(x)
    print(x)
    # 定义要求导/梯度的函数
    with tf.GradientTape() as tape:  # 追踪梯度
        s = tf.nn.leaky_relu(x, alpha=0.1)

    # 计算 函数 在 x 位置的导数值
    grad = tape.gradient(s, x)  # 计算梯度
    print(grad)

    plt.plot(x, y, color='r', label="relu")
    plt.plot(x, grad, color='y', label='tanh gradient')
    # 设置图例
    plt.legend(loc=2)
    plt.title("leaky_relu(x) and leaky_relu(x) gradient")
    plt.grid()
    plt.xlabel(xlabel='x')
    plt.ylabel(ylabel='y')
    plt.xlim(-10, 10)
    # plt.ylim(-1, 1)
    plt.show()